import numpy as np
import pandas as pd
from scipy.linalg import eigh
import matplotlib.pyplot as plt

def calculate_rmse(estimated_doa, real_doa):
    """
    计算估计的 DOA 与真实 DOA 之间的 RMSE。
    
    参数：
        estimated_doa (list or numpy array): 估计的 DOA 值列表。
        real_doa (list or numpy array): 真实的 DOA 值列表。
        
    返回：
        rmse (float): 估计 DOA 与真实 DOA 之间的均方根误差。
    """
    # 确保两个列表长度相等
    assert len(estimated_doa) == len(real_doa), "估计的 DOA 和真实 DOA 长度必须相等"
    
    # 计算 RMSE
    squared_errors = [(e - r) ** 2 for e, r in zip(estimated_doa, real_doa)]
    rmse = np.sqrt(np.mean(squared_errors))
    
    return rmse

def calculate_rmse_complex(estimated_doa, real_doa):
    """
    计算虚数估计 DOA 与真实 DOA 之间的 RMSE。
    """
    # 确保两个列表长度相等
    assert len(estimated_doa) == len(real_doa), "估计的 DOA 和真实 DOA 长度必须相等"
    
    # 分别计算实部和虚部的误差
    real_part_error = [(e.real - r.real) ** 2 for e, r in zip(estimated_doa, real_doa)]
    imag_part_error = [(e.imag - r.imag) ** 2 for e, r in zip(estimated_doa, real_doa)]
    
    # 计算总的 RMSE
    rmse = np.sqrt(np.mean(real_part_error + imag_part_error))
    
    return rmse

def generate_DOAs(num_doa):
    """
    生成指定数量的随机 DOA，每个 DOA 都是 10 的整数倍
    """
    # 设置 DOA 的范围，例如 -90 到 90 度
    min_doa = -90
    max_doa = 90
    
    # 生成范围内的所有 10 的整数倍
    possible_DOAs = np.arange(min_doa, max_doa + 1, 10)
    
    # 从这些可能的值中随机选择
    DOAs = np.random.choice(possible_DOAs, size=num_doa, replace=True)
    
    return DOAs

def steering_vector(M, angle, wavelength=1, d=0.5):
    theta = np.deg2rad(angle)
    return np.exp(1j * 2 * np.pi * d * np.outer(np.arange(M) , np.sin(theta) / wavelength))

def generate_mutual_coupling_matrix(M,qu):
    """
    生成一个 MxM 的互耦矩阵
    对角线元素为 1，非对角线元素随着距离增加快速衰减，并且是复数
    """
    C = np.eye(M, dtype=complex)  # 初始化为单位矩阵

    for i in range(M):
        for j in range(i + 1, M):
            distance = abs(i - j)
            real_part = np.exp(-distance / 2)  # 实部快速衰减
            imag_part = np.exp(-distance / 2)  # 虚部快速衰减
            C[i, j] = real_part + 1j * imag_part
            C[j, i] = C[i, j]  # 对称性
    
    # m2 = M-qu
    # for i in range(m2,M,1):
    #     for j in range(m2):
    #         if j>0 and i>=m2+j:
    #             C[i,j] = 0
    #             C[j,i] = 0
    #         elif j==0:
    #             C[i,j]=0
    #             C[j,i]=0
    # 设置左下角和右上角部分为零
    for i in range(M):
        for j in range(M):
            if (j >= M - qu and i <= j - (M - qu)) or (i >= M - qu and j <= i - (M - qu)):
                C[i, j] = 0

    return C

def generate_small_perturbation_matrix(M, epsilon=1e-6):
    """
    生成一个 MxM 的小扰动矩阵
    所有元素为非常小的复数，用于加到原来的矩阵上跳出局部最优解
    """
    # 初始化矩阵 P 为 MxM 的零矩阵
    P = np.zeros((M, M), dtype=complex)

    # 生成小扰动复数元素
    for i in range(M):
        for j in range(M):
            real_part = epsilon * np.random.rand()  # 小实部
            imag_part = epsilon * np.random.rand()  # 小虚部
            complex_value = real_part + 1j * imag_part  # 生成小复数
            P[i, j] = complex_value

    return P

def generate_complex_matrix(M):
    """
    生成 MxM 的矩阵 C
    对角线元素为 1，其他非对角线元素为复数，并且所有斜对角线元素相同
    """
    # 初始化矩阵 C 为 MxM 的单位矩阵
    C = np.eye(M, dtype=complex)

    # 生成复数元素
    for i in range(1, M):
        real_part = np.random.rand()  # 随机实部
        imag_part = np.random.rand()  # 随机虚部
        complex_value = real_part + 1j * imag_part  # 生成复数

        for j in range(M - i):
            C[j, j + i] = complex_value  # 上三角斜对角线
            C[j + i, j] = complex_value  # 下三角斜对角线
    return C

def awgn(signal, snr_db):
    snr_linear = 10 ** (snr_db / 10.0)
    power_signal = np.sum(np.abs(signal) ** 2) / signal.size
    power_noise = power_signal / snr_linear
    noise = np.sqrt(power_noise / 2) * (np.random.randn(*signal.shape) + 1j * np.random.randn(*signal.shape))
    return signal + noise

def generate_mat_J(mat_C, q, M):
    mat_J = np.zeros((M, M))#, #dtype=complex)
    for i in range(M):
        for j in range(M):
            if mat_C[i, j] == q:
                mat_J[i, j] = 1
    return mat_J

def gMat_J_all(M,mat_C,Q):
    Mat_all_J = np.zeros((M, M * Q))
    rows = mat_C.shape[0]
    for q in range(Q):
        mat_J = generate_mat_J(mat_C, q + 1, rows)
        # 将 mat_J 按列堆叠到 Mat_all_J 中
        start_col = q * M
        end_col = (q + 1) * M
        Mat_all_J[:, start_col:end_col] = mat_J
    return Mat_all_J

def gMat_T(M, DoA_i, mat_C,Q_len):
    d = 0.5  # number of target
    Geo_Tx = np.arange(M)  # geometry of Tx
    rows = mat_C.shape[0]
    Q = Q_len
    Mat_T = np.zeros((M, Q), dtype=complex)
    for q in range(Q):
        mat_J = generate_mat_J(mat_C, q + 1, rows)  # MATLAB 1-indexed, Python 0-indexed
        At = np.exp(1j * 2 * np.pi * d * Geo_Tx * np.sin(np.deg2rad(DoA_i)))
        Jq = mat_J @ At
        Mat_T[:, q] = Jq
    return Mat_T

# MC互耦阵列的索引
def gMatC(size):
    mat = np.zeros((size, size), dtype=float)
    last_s = 1
    for i in range(size):
        for j in range(size):
            if i==j:
                mat[i,j]=1
            #elif j<i:
            else:
                mat[i, j] = 1+abs(i - j)
    return mat

def gSig(M,Doa,L,MCM,K=1,SNR=60):
    a = steering_vector(M,Doa)
    S = np.random.randn(K, L) + 1j * np.random.randn(K, L)  # 信号矩阵    
    sig = MCM @ a @ S # 生成信号
    X = awgn(sig, SNR)  # 添加噪声后信号
    R = (X @ X.conj().T) / L  # 信号协方差矩阵
    D, En = eigh(R, subset_by_index=[0, M-K-1]) # 噪声子空间与噪声特征值
    return En
    
# MUSIC_test(M,L,60)    
def MUSIC_test(M,L,DOA):
    En = gSig(M,DOA,L,np.eye(M,M))
    search_grid = np.arange(-90, 90, 1)  # 搜索范围
    G = np.zeros(len(search_grid))
    for idx,theta in enumerate(search_grid):
        a = steering_vector(M,theta)
        G[idx] = 10*np.log10(1/abs(a.conj().T @ En @ En.conj().T @ a))
    plt.figure()
    plt.plot(search_grid,G)
    plt.show()    

tflag = 0
def Pmu(En_lst,MC_gu,M):
    global tflag
    P = 0 # 越大越好
    Doa_s = []
    search_grid = np.arange(-90, 90, 1)  # 搜索范围
    for En in (En_lst):
        gMax = 0
        tdoa = -999
        g_lst = []
        for idx,theta in enumerate(search_grid):
            a = steering_vector(M,theta)
            g = (1/abs(a.conj().T @ MC_gu.conj().T @ En @ En.conj().T @ MC_gu @ a))
            g_lst.append(g[0])
            if gMax<g[0]:
                gMax = g[0]
                tdoa = theta
        # if tflag==1:
        # plt.figure()
        # plt.plot(search_grid,g_lst)
        # plt.show()
        P += gMax
        Doa_s.append(tdoa)
    
    tflag = 1
    return Doa_s,P
   
def updateC2(phi,M,w,mat_C,En_lst,Q_len,Doa_zhi,doa_len):
    Q_sum = 0
    for idx in range(doa_len):
        En = En_lst[idx]
        Mat_T = gMat_T(M, Doa_zhi, mat_C,Q_len)
        Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
        Q_sum += Q_mat
    c = (np.linalg.inv(Q_sum) @ w) /(w.conj().T @ Q_sum @ w)
    c = c/c[0]
    cost_val2 = abs(c.conj().T @ Q_sum @ c)
    psi = np.kron(c,np.eye(M))
    C = phi @ psi
    
    return C,cost_val2  
    
def updateC(phi,M,w,mat_C,En_lst,doa_Gu_lst,Q_len):
    Q_sum = 0
    for idx,doai in enumerate(doa_Gu_lst):
        En = En_lst[idx]
        Mat_T = gMat_T(M, doai, mat_C,Q_len)
        Q_mat = Mat_T.conj().T @ (En @ En.conj().T) @ Mat_T
        Q_sum += Q_mat
    c = (np.linalg.inv(Q_sum) @ w) /(w.conj().T @ Q_sum @ w)
    c = c/c[0]
    cost_val2 = abs(c.conj().T @ Q_sum @ c) # 越小越好
    psi = np.kron(c,np.eye(M))
    C = phi @ psi
    
    return C,cost_val2

def main():
    M = 8
    lambdaf = 1
    d = 0.5
    L = 300
    qu = 4
    MCM = generate_mutual_coupling_matrix(M,qu)
    mcm_df = pd.DataFrame(MCM)
    #DOAs = np.full(1, 50)
    DOAs = [50,80]
    print("预设角度:",DOAs)
    doaLen = len(DOAs)
    En_lst = []
    snr = 20
    for doa in DOAs:
        En = gSig(M,doa,L,MCM,1,snr)
        En_lst.append(En)
    
    mat_C = gMatC(M) # 初始化 MC 下标矩阵
    Q_len = M-qu
    w = np.eye(Q_len, 1)
    phi = gMat_J_all(M,mat_C,Q_len)
    
    # 初始化MCM
    #C1 = generate_complex_matrix(M)
    identity_matrix  = (np.eye(M))
    C1 = identity_matrix
    tgC = ""
    tgVal = 999
    search_grid = np.arange(-90, 90, 1)  # 搜索范围
    search_grid_r = np.zeros(len(search_grid))
    rightDoa = -999
    
    for i in range(5):
        doa_Gu_lst,P_cost = Pmu(En_lst,C1,M)
        print(f"估计角度:{doa_Gu_lst[0]}")
        C,cost_val2 = updateC(phi,M,w,mat_C,En_lst,doa_Gu_lst,Q_len)
        print("代价函数值：",cost_val2)
        C1 = C
        
    Ori_c = mcm_df.iloc[:, 0].tolist()
    guji_c = pd.DataFrame(C).iloc[:,0].tolist()
    print(' '.join(map(str, Ori_c)))
    print(' '.join(map(str, guji_c)))
    rmse_i = calculate_rmse_complex(Ori_c,guji_c)
    print("RMSE:",rmse_i)


if __name__=="__main__":
    main()